"""
数据源服务
处理数据源的CRUD操作、连接测试、数据预览等业务逻辑
"""

import json
import os
from typing import List, Optional, Tuple, Dict, Any
from sqlalchemy.orm import Session
from sqlalchemy import func
from fastapi import HTTPException, status, UploadFile
import pymysql
import pandas as pd
import tempfile

from models.data_source import DataSource, DataSourceType
from schemas.data_source import (
    DataSourceCreate, DataSourceUpdate, DataSourceResponse, DataSourceList,
    ConnectionTestRequest, ConnectionTestResponse, DataPreviewRequest, DataPreviewResponse,
    MySQLConfig, CSVConfig, ExcelConfig
)
from config import settings


class DataSourceService:
    """数据源服务类"""
    
    @staticmethod
    def create_data_source(db: Session, data_source_data: DataSourceCreate, user_id: int) -> DataSource:
        """创建数据源"""
        # 检查名称是否已存在
        existing = db.query(DataSource).filter(
            DataSource.name == data_source_data.name,
            DataSource.user_id == user_id
        ).first()
        
        if existing:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="数据源名称已存在"
            )
        
        # 创建数据源
        db_data_source = DataSource(
            name=data_source_data.name,
            type=data_source_data.type,
            config=json.dumps(data_source_data.config, ensure_ascii=False),
            description=data_source_data.description,
            user_id=user_id
        )
        
        db.add(db_data_source)
        db.commit()
        db.refresh(db_data_source)
        return db_data_source
    
    @staticmethod
    def get_data_sources(
        db: Session, 
        user_id: int, 
        page: int = 1, 
        size: int = 10,
        search: Optional[str] = None,
        data_type: Optional[DataSourceType] = None
    ) -> DataSourceList:
        """获取数据源列表"""
        query = db.query(DataSource).filter(DataSource.user_id == user_id)
        
        # 搜索过滤
        if search:
            query = query.filter(DataSource.name.contains(search))
        
        # 类型过滤
        if data_type:
            query = query.filter(DataSource.type == data_type)
        
        # 总数
        total = query.count()
        
        # 分页
        items = query.offset((page - 1) * size).limit(size).all()
        
        # 计算总页数
        pages = (total + size - 1) // size
        
        return DataSourceList(
            items=[DataSourceResponse.from_orm(item) for item in items],
            total=total,
            page=page,
            size=size,
            pages=pages
        )
    
    @staticmethod
    def get_data_source(db: Session, data_source_id: int, user_id: int) -> Optional[DataSource]:
        """获取单个数据源"""
        return db.query(DataSource).filter(
            DataSource.id == data_source_id,
            DataSource.user_id == user_id
        ).first()
    
    @staticmethod
    def update_data_source(
        db: Session, 
        data_source_id: int, 
        data_source_data: DataSourceUpdate, 
        user_id: int
    ) -> Optional[DataSource]:
        """更新数据源"""
        db_data_source = DataSourceService.get_data_source(db, data_source_id, user_id)
        if not db_data_source:
            return None
        
        # 检查名称是否已存在（排除自己）
        if data_source_data.name:
            existing = db.query(DataSource).filter(
                DataSource.name == data_source_data.name,
                DataSource.user_id == user_id,
                DataSource.id != data_source_id
            ).first()
            
            if existing:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="数据源名称已存在"
                )
        
        # 更新字段
        update_data = data_source_data.dict(exclude_unset=True)
        for field, value in update_data.items():
            if field == 'config':
                setattr(db_data_source, field, json.dumps(value, ensure_ascii=False))
            else:
                setattr(db_data_source, field, value)
        
        db.commit()
        db.refresh(db_data_source)
        return db_data_source
    
    @staticmethod
    def delete_data_source(db: Session, data_source_id: int, user_id: int) -> bool:
        """删除数据源"""
        db_data_source = DataSourceService.get_data_source(db, data_source_id, user_id)
        if not db_data_source:
            return False
        
        # 如果是文件类型的数据源，删除对应的文件
        if db_data_source.type in [DataSourceType.CSV, DataSourceType.EXCEL]:
            try:
                config = json.loads(db_data_source.config)
                file_path = config.get('file_path')
                if file_path and os.path.exists(file_path):
                    os.remove(file_path)
            except Exception:
                pass  # 忽略文件删除错误
        
        db.delete(db_data_source)
        db.commit()
        return True

    @staticmethod
    def test_connection(test_request: ConnectionTestRequest) -> ConnectionTestResponse:
        """测试数据源连接"""
        try:
            if test_request.type == DataSourceType.MYSQL:
                return DataSourceService._test_mysql_connection(test_request.config)
            elif test_request.type == DataSourceType.CSV:
                return DataSourceService._test_csv_connection(test_request.config)
            elif test_request.type == DataSourceType.EXCEL:
                return DataSourceService._test_excel_connection(test_request.config)
            else:
                return ConnectionTestResponse(
                    success=False,
                    message="不支持的数据源类型"
                )
        except Exception as e:
            return ConnectionTestResponse(
                success=False,
                message=f"连接测试失败: {str(e)}"
            )

    @staticmethod
    def _test_mysql_connection(config: Dict[str, Any]) -> ConnectionTestResponse:
        """测试MySQL连接"""
        try:
            mysql_config = MySQLConfig(**config)
            connection = pymysql.connect(
                host=mysql_config.host,
                port=mysql_config.port,
                user=mysql_config.username,
                password=mysql_config.password,
                database=mysql_config.database,
                charset=mysql_config.charset,
                connect_timeout=10
            )

            with connection.cursor() as cursor:
                cursor.execute("SELECT 1")
                cursor.fetchone()

            connection.close()

            return ConnectionTestResponse(
                success=True,
                message="MySQL连接成功",
                details={
                    "host": mysql_config.host,
                    "port": mysql_config.port,
                    "database": mysql_config.database
                }
            )
        except Exception as e:
            return ConnectionTestResponse(
                success=False,
                message=f"MySQL连接失败: {str(e)}"
            )

    @staticmethod
    def _test_csv_connection(config: Dict[str, Any]) -> ConnectionTestResponse:
        """测试CSV文件连接"""
        try:
            csv_config = CSVConfig(**config)

            if not os.path.exists(csv_config.file_path):
                return ConnectionTestResponse(
                    success=False,
                    message="CSV文件不存在"
                )

            # 尝试读取前几行
            df = pd.read_csv(
                csv_config.file_path,
                encoding=csv_config.encoding,
                delimiter=csv_config.delimiter,
                header=0 if csv_config.has_header else None,
                nrows=5
            )

            return ConnectionTestResponse(
                success=True,
                message="CSV文件读取成功",
                details={
                    "file_path": csv_config.file_path,
                    "rows": len(df),
                    "columns": len(df.columns),
                    "encoding": csv_config.encoding
                }
            )
        except Exception as e:
            return ConnectionTestResponse(
                success=False,
                message=f"CSV文件读取失败: {str(e)}"
            )

    @staticmethod
    def _test_excel_connection(config: Dict[str, Any]) -> ConnectionTestResponse:
        """测试Excel文件连接"""
        try:
            excel_config = ExcelConfig(**config)

            if not os.path.exists(excel_config.file_path):
                return ConnectionTestResponse(
                    success=False,
                    message="Excel文件不存在"
                )

            # 尝试读取前几行
            df = pd.read_excel(
                excel_config.file_path,
                sheet_name=excel_config.sheet_name,
                header=0 if excel_config.has_header else None,
                skiprows=excel_config.start_row,
                nrows=5
            )

            return ConnectionTestResponse(
                success=True,
                message="Excel文件读取成功",
                details={
                    "file_path": excel_config.file_path,
                    "sheet_name": excel_config.sheet_name or "第一个工作表",
                    "rows": len(df),
                    "columns": len(df.columns)
                }
            )
        except Exception as e:
            return ConnectionTestResponse(
                success=False,
                message=f"Excel文件读取失败: {str(e)}"
            )

    @staticmethod
    def get_data_source_tables(db: Session, data_source_id: int, user_id: int):
        """获取数据源表列表"""
        # 获取数据源
        data_source = DataSourceService.get_data_source(db, data_source_id, user_id)
        if not data_source:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="数据源不存在"
            )

        try:
            config = json.loads(data_source.config)

            if data_source.type == DataSourceType.MYSQL:
                return DataSourceService._get_mysql_tables(config)
            elif data_source.type == DataSourceType.CSV:
                # CSV文件没有表的概念，返回文件信息
                return {
                    "tables": [{"name": "csv_data", "comment": "CSV文件数据", "row_count": None}],
                    "total_count": 1
                }
            elif data_source.type == DataSourceType.EXCEL:
                return DataSourceService._get_excel_sheets(config)
            else:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="不支持的数据源类型"
                )
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"获取表列表失败: {str(e)}"
            )

    @staticmethod
    def preview_data(db: Session, preview_request: DataPreviewRequest, user_id: int) -> DataPreviewResponse:
        """预览数据源数据"""
        # 获取数据源
        data_source = DataSourceService.get_data_source(db, preview_request.data_source_id, user_id)
        if not data_source:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="数据源不存在"
            )

        try:
            config = json.loads(data_source.config)

            if data_source.type == DataSourceType.MYSQL:
                return DataSourceService._preview_mysql_data(config, preview_request.limit, preview_request.table_name)
            elif data_source.type == DataSourceType.CSV:
                return DataSourceService._preview_csv_data(config, preview_request.limit)
            elif data_source.type == DataSourceType.EXCEL:
                return DataSourceService._preview_excel_data(config, preview_request.limit, preview_request.table_name)
            else:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="不支持的数据源类型"
                )
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"数据预览失败: {str(e)}"
            )

    @staticmethod
    def _get_mysql_tables(config: Dict[str, Any]):
        """获取MySQL数据库表列表"""
        mysql_config = MySQLConfig(**config)
        connection = pymysql.connect(
            host=mysql_config.host,
            port=mysql_config.port,
            user=mysql_config.username,
            password=mysql_config.password,
            database=mysql_config.database,
            charset=mysql_config.charset
        )

        try:
            with connection.cursor() as cursor:
                # 获取表信息
                cursor.execute("""
                    SELECT
                        TABLE_NAME,
                        TABLE_COMMENT,
                        TABLE_ROWS
                    FROM
                        INFORMATION_SCHEMA.TABLES
                    WHERE
                        TABLE_SCHEMA = %s
                        AND TABLE_TYPE = 'BASE TABLE'
                    ORDER BY TABLE_NAME
                """, (mysql_config.database,))

                tables = []
                for row in cursor.fetchall():
                    tables.append({
                        "name": row[0],
                        "comment": row[1] if row[1] else None,
                        "row_count": row[2] if row[2] else None
                    })

                return {
                    "tables": tables,
                    "total_count": len(tables)
                }
        finally:
            connection.close()

    @staticmethod
    def _get_excel_sheets(config: Dict[str, Any]):
        """获取Excel文件工作表列表"""
        excel_config = ExcelConfig(**config)

        try:
            # 读取Excel文件的所有工作表名称
            excel_file = pd.ExcelFile(excel_config.file_path)
            sheets = []

            for sheet_name in excel_file.sheet_names:
                # 获取每个工作表的行数
                df = pd.read_excel(excel_config.file_path, sheet_name=sheet_name, nrows=0)
                total_df = pd.read_excel(excel_config.file_path, sheet_name=sheet_name)

                sheets.append({
                    "name": sheet_name,
                    "comment": f"Excel工作表: {sheet_name}",
                    "row_count": len(total_df)
                })

            return {
                "tables": sheets,
                "total_count": len(sheets)
            }
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"读取Excel文件失败: {str(e)}"
            )

    @staticmethod
    def _preview_mysql_data(config: Dict[str, Any], limit: int, table_name: Optional[str] = None) -> DataPreviewResponse:
        """预览MySQL数据"""
        mysql_config = MySQLConfig(**config)
        connection = pymysql.connect(
            host=mysql_config.host,
            port=mysql_config.port,
            user=mysql_config.username,
            password=mysql_config.password,
            database=mysql_config.database,
            charset=mysql_config.charset
        )

        try:
            with connection.cursor() as cursor:
                # 如果没有指定表名，获取第一个表
                if not table_name:
                    cursor.execute("SHOW TABLES")
                    tables = [row[0] for row in cursor.fetchall()]

                    if not tables:
                        return DataPreviewResponse(
                            columns=[],
                            data=[],
                            total_rows=0,
                            sample_rows=0,
                            table_name=None
                        )

                    table_name = tables[0]

                # 获取总行数
                cursor.execute(f"SELECT COUNT(*) FROM `{table_name}`")
                total_rows = cursor.fetchone()[0]

                # 获取列信息
                cursor.execute(f"DESCRIBE `{table_name}`")
                columns = [row[0] for row in cursor.fetchall()]

                # 获取数据
                cursor.execute(f"SELECT * FROM `{table_name}` LIMIT {limit}")
                rows = cursor.fetchall()

                # 转换数据格式
                data = []
                for row in rows:
                    data.append([str(cell) if cell is not None else "" for cell in row])

                return DataPreviewResponse(
                    columns=columns,
                    data=data,
                    total_rows=total_rows,
                    sample_rows=len(data),
                    table_name=table_name
                )
        finally:
            connection.close()

    @staticmethod
    def _preview_csv_data(config: Dict[str, Any], limit: int) -> DataPreviewResponse:
        """预览CSV数据"""
        csv_config = CSVConfig(**config)

        # 读取数据
        df = pd.read_csv(
            csv_config.file_path,
            encoding=csv_config.encoding,
            delimiter=csv_config.delimiter,
            header=0 if csv_config.has_header else None,
            nrows=limit
        )

        # 获取总行数（重新读取文件）
        total_df = pd.read_csv(
            csv_config.file_path,
            encoding=csv_config.encoding,
            delimiter=csv_config.delimiter,
            header=0 if csv_config.has_header else None
        )
        total_rows = len(total_df)

        # 处理列名
        if csv_config.has_header:
            columns = df.columns.tolist()
        else:
            columns = [f"Column_{i+1}" for i in range(len(df.columns))]
            df.columns = columns

        # 转换数据
        data = []
        for _, row in df.iterrows():
            data.append([str(cell) if pd.notna(cell) else "" for cell in row])

        return DataPreviewResponse(
            columns=columns,
            data=data,
            total_rows=total_rows,
            sample_rows=len(data)
        )

    @staticmethod
    def _preview_excel_data(config: Dict[str, Any], limit: int, sheet_name: Optional[str] = None) -> DataPreviewResponse:
        """预览Excel数据"""
        excel_config = ExcelConfig(**config)

        # 确定要使用的工作表名称
        target_sheet = sheet_name or excel_config.sheet_name

        # 读取数据
        df = pd.read_excel(
            excel_config.file_path,
            sheet_name=target_sheet,
            header=0 if excel_config.has_header else None,
            skiprows=excel_config.start_row,
            nrows=limit
        )

        # 获取总行数（重新读取文件）
        total_df = pd.read_excel(
            excel_config.file_path,
            sheet_name=excel_config.sheet_name,
            header=0 if excel_config.has_header else None,
            skiprows=excel_config.start_row
        )
        total_rows = len(total_df)

        # 处理列名
        if excel_config.has_header:
            columns = df.columns.tolist()
        else:
            columns = [f"Column_{i+1}" for i in range(len(df.columns))]
            df.columns = columns

        # 转换数据
        data = []
        for _, row in df.iterrows():
            data.append([str(cell) if pd.notna(cell) else "" for cell in row])

        return DataPreviewResponse(
            columns=columns,
            data=data,
            total_rows=total_rows,
            sample_rows=len(data)
        )

    @staticmethod
    def upload_file(file: UploadFile, user_id: int) -> str:
        """上传文件并返回文件路径"""
        # 检查文件类型
        allowed_extensions = ['.csv', '.xlsx', '.xls']
        file_extension = os.path.splitext(file.filename)[1].lower()

        if file_extension not in allowed_extensions:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"不支持的文件类型，仅支持: {', '.join(allowed_extensions)}"
            )

        # 检查文件大小（限制为10MB）
        if file.size > settings.max_file_size:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"文件大小超过限制（{settings.max_file_size / 1024 / 1024:.1f}MB）"
            )

        # 创建用户专属目录
        user_upload_dir = os.path.join(settings.upload_dir, f"user_{user_id}")
        os.makedirs(user_upload_dir, exist_ok=True)

        # 生成唯一文件名
        import uuid
        unique_filename = f"{uuid.uuid4().hex}_{file.filename}"
        file_path = os.path.join(user_upload_dir, unique_filename)

        # 保存文件
        try:
            with open(file_path, "wb") as buffer:
                content = file.file.read()
                buffer.write(content)
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"文件保存失败: {str(e)}"
            )

        return file_path

    @staticmethod
    def get_excel_sheets(file_path: str) -> List[str]:
        """获取Excel文件的所有工作表名称"""
        try:
            excel_file = pd.ExcelFile(file_path)
            return excel_file.sheet_names
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"读取Excel文件失败: {str(e)}"
            )

    @staticmethod
    def get_data_source_statistics(db: Session, user_id: int) -> Dict[str, Any]:
        """获取数据源统计信息"""
        # 总数统计
        total = db.query(DataSource).filter(DataSource.user_id == user_id).count()

        # 按类型统计
        type_stats = db.query(
            DataSource.type,
            func.count(DataSource.id).label('count')
        ).filter(DataSource.user_id == user_id).group_by(DataSource.type).all()

        type_counts = {item.type: item.count for item in type_stats}

        return {
            "total": total,
            "mysql": type_counts.get(DataSourceType.MYSQL, 0),
            "csv": type_counts.get(DataSourceType.CSV, 0),
            "excel": type_counts.get(DataSourceType.EXCEL, 0)
        }

    @staticmethod
    def get_table_fields(db: Session, data_source_id: int, table_name: str, user_id: int, identify_types: bool = False) -> List[Dict[str, Any]]:
        """获取表字段信息"""
        data_source = db.query(DataSource).filter(
            DataSource.id == data_source_id,
            DataSource.user_id == user_id
        ).first()

        if not data_source:
            raise HTTPException(status_code=404, detail="数据源不存在")

        if data_source.type == DataSourceType.MYSQL:
            return DataSourceService._get_mysql_table_fields(data_source, table_name, identify_types)
        else:
            raise HTTPException(status_code=400, detail="暂不支持该数据源类型的字段获取")

    @staticmethod
    def get_table_data(db: Session, data_source_id: int, table_name: str, fields: str, limit: int, conditions: str, user_id: int) -> List[Dict[str, Any]]:
        """获取表数据"""
        data_source = db.query(DataSource).filter(
            DataSource.id == data_source_id,
            DataSource.user_id == user_id
        ).first()

        if not data_source:
            raise HTTPException(status_code=404, detail="数据源不存在")

        if data_source.type == DataSourceType.MYSQL:
            return DataSourceService._get_mysql_table_data(data_source, table_name, fields, limit, conditions)
        else:
            raise HTTPException(status_code=400, detail="暂不支持该数据源类型的数据获取")

    @staticmethod
    def _get_mysql_table_fields(data_source: DataSource, table_name: str, identify_types: bool = False) -> List[Dict[str, Any]]:
        """获取MySQL表字段"""
        try:
            config = json.loads(data_source.config) if isinstance(data_source.config, str) else data_source.config

            connection = pymysql.connect(
                host=config.get('host'),
                port=config.get('port', 3306),
                user=config.get('username'),
                password=config.get('password'),
                database=config.get('database'),
                charset=config.get('charset', 'utf8mb4')
            )

            fields = []
            with connection.cursor() as cursor:
                cursor.execute(f"DESCRIBE {table_name}")
                columns = cursor.fetchall()

                for column in columns:
                    field_name = column[0]
                    field_type = column[1]

                    field_info = {
                        "name": field_name,
                        "type": field_type,
                        "comment": "",
                    }

                    if identify_types:
                        # 简单的类型识别
                        if 'int' in field_type.lower() or 'decimal' in field_type.lower() or 'float' in field_type.lower():
                            field_info["field_type"] = "number"
                        elif 'date' in field_type.lower() or 'time' in field_type.lower():
                            field_info["field_type"] = "date"
                        else:
                            field_info["field_type"] = "string"

                        # 识别邮箱和姓名字段
                        field_info["is_email"] = 'email' in field_name.lower() or 'mail' in field_name.lower()
                        field_info["is_name"] = 'name' in field_name.lower() or '姓名' in field_name.lower()

                    fields.append(field_info)

            connection.close()
            return fields

        except Exception as e:
            raise HTTPException(status_code=500, detail=f"获取字段失败: {str(e)}")

    @staticmethod
    def _get_mysql_table_data(data_source: DataSource, table_name: str, fields: str, limit: int, conditions: str) -> List[Dict[str, Any]]:
        """获取MySQL表数据"""
        try:
            config = json.loads(data_source.config) if isinstance(data_source.config, str) else data_source.config

            connection = pymysql.connect(
                host=config.get('host'),
                port=config.get('port', 3306),
                user=config.get('username'),
                password=config.get('password'),
                database=config.get('database'),
                charset=config.get('charset', 'utf8mb4')
            )

            # 构建SQL查询
            field_list = [f.strip() for f in fields.split(',')]
            sql = f"SELECT {', '.join(field_list)} FROM {table_name}"

            # 添加筛选条件
            if conditions:
                try:
                    condition_list = json.loads(conditions)
                    if condition_list:
                        where_clauses = []
                        for condition in condition_list:
                            field = condition.get('field')
                            operator = condition.get('operator', '=')
                            value = condition.get('value')

                            if field and value:
                                if operator == 'LIKE':
                                    where_clauses.append(f"{field} LIKE '%{value}%'")
                                else:
                                    where_clauses.append(f"{field} {operator} '{value}'")

                        if where_clauses:
                            sql += " WHERE " + " AND ".join(where_clauses)
                except:
                    pass  # 忽略条件解析错误

            sql += f" LIMIT {limit}"

            data = []
            with connection.cursor() as cursor:
                cursor.execute(sql)
                columns = [desc[0] for desc in cursor.description]
                rows = cursor.fetchall()

                for row in rows:
                    row_dict = {}
                    for i, value in enumerate(row):
                        row_dict[columns[i]] = value
                    data.append(row_dict)

            connection.close()
            return data

        except Exception as e:
            raise HTTPException(status_code=500, detail=f"获取数据失败: {str(e)}")
